#include <stdio.h>
#include <stdlib.h>
#include <netinet/in.h>
#include <tss2/tss2_rc.h>
#include <tss2/tss2_sys.h>
#include <tss2/tss2_tctildr.h>

#ifndef TCTI_NAME
#define TCTI_NAME "mssim"
#endif

#define CHECK_RC(rc, message, handle) \
    do { \
        if ((rc) != TSS2_RC_SUCCESS) { \
            fprintf(stderr, "错误: %s 失败! TSS2_RC: 0x%x, %s\n", \
                    (message), (rc), Tss2_RC_Decode(rc)); \
            goto cleanup; \
        } \
    } while (0)

TSS2_RC init_tcti_context(TSS2_TCTI_CONTEXT **tcti_context);
TSS2_RC init_sys_context(TSS2_TCTI_CONTEXT *tcti_context, TSS2_SYS_CONTEXT **sys_context);
TSS2_RC startup_tpm(TSS2_SYS_CONTEXT *sys_context);
TSS2_RC get_tpm_mfr_id(TSS2_SYS_CONTEXT *sys_context, UINT32 *manufacturer_id);
void clean(TSS2_TCTI_CONTEXT *tcti_context, TSS2_SYS_CONTEXT *sys_context);

int main() {
    TSS2_RC rc = TSS2_RC_SUCCESS;
    TSS2_TCTI_CONTEXT *tcti_context = NULL;
    TSS2_SYS_CONTEXT *sys_context = NULL;
    UINT32 mfr_id = 0;

    /* --- 初始化 TCTI Context --- */
    rc = init_tcti_context(&tcti_context);
    CHECK_RC(rc, "初始化 TCTI Context", tcti_context);

    /* --- 初始化 SAPI Context --- */
    rc = init_sys_context(tcti_context, &sys_context);
    CHECK_RC(rc, "初始化 SAPI Context", sys_context);

    /* --- 启动 TPM --- */
    rc = startup_tpm(sys_context);
    CHECK_RC(rc, "启动 TPM", sys_context);

    /* --- 查询 TPM 制造商 ID --- */
    rc = get_tpm_mfr_id(sys_context, &mfr_id);
    CHECK_RC(rc, "查询 TPM 制造商 ID", sys_context);
    // 将网络字节序转换为主机字节序
    UINT32 mfr_id_host = ntohl(mfr_id);
    const char* p = (char*)&(mfr_id_host);
    printf("TPM 制造商: %c%c%c%c\n", p[0], p[1], p[2], p[3]);

    /* --- 清理资源 --- */
cleanup:
    clean(tcti_context, sys_context);
    
    return !(rc == TSS2_RC_SUCCESS); 
}

// 初始化 TCTI Context
TSS2_RC init_tcti_context(TSS2_TCTI_CONTEXT **tcti_context) {
    return Tss2_TctiLdr_Initialize(TCTI_NAME, tcti_context);
}

// 初始化 SAPI Context
TSS2_RC init_sys_context(TSS2_TCTI_CONTEXT *tcti_context, TSS2_SYS_CONTEXT **sys_context) {
    if (tcti_context == NULL) {
        return TSS2_SYS_RC_BAD_REFERENCE;
    }

    size_t sys_context_size = Tss2_Sys_GetContextSize(0);
    *sys_context = (TSS2_SYS_CONTEXT *)malloc(sys_context_size);
    if (*sys_context == NULL) {
        return TSS2_BASE_RC_MEMORY;
    }
    
    TSS2_ABI_VERSION abi_version = TSS2_ABI_VERSION_CURRENT;
    TSS2_RC rc = Tss2_Sys_Initialize(*sys_context, sys_context_size, tcti_context, &abi_version);
    if (rc != TSS2_RC_SUCCESS) {
        free(*sys_context);
        *sys_context = NULL;
    }
    
    return rc;
}

// 启动 TPM
TSS2_RC startup_tpm(TSS2_SYS_CONTEXT *sys_context) {
    if (sys_context == NULL) {
        return TSS2_SYS_RC_BAD_REFERENCE;
    }
    
    TSS2_RC rc = Tss2_Sys_GetCapability(sys_context, NULL, 
                                        TPM2_CAP_TPM_PROPERTIES, TPM2_PT_MANUFACTURER, 
                                        1, NULL, NULL, NULL);
    if (rc == TPM2_RC_INITIALIZE) {
        rc = Tss2_Sys_Startup(sys_context, TPM2_SU_CLEAR);
    } 
    
    return rc;
}

// 查询 TPM 制造商 ID
TSS2_RC get_tpm_mfr_id(TSS2_SYS_CONTEXT *sys_context, UINT32 *mfr_id) {

    if (sys_context == NULL || mfr_id == NULL) {
        return TSS2_SYS_RC_BAD_REFERENCE;
    }
    
    TPMS_CAPABILITY_DATA cap_data;
    TSS2_RC rc = Tss2_Sys_GetCapability(sys_context, NULL, 
                                        TPM2_CAP_TPM_PROPERTIES, TPM2_PT_MANUFACTURER,
                                        1, NULL, &cap_data, NULL);
    
    if (rc == TSS2_RC_SUCCESS) {
        *mfr_id = cap_data.data.tpmProperties.tpmProperty[0].value;
    }
    
    return rc;
}

// 清理资源
void clean(TSS2_TCTI_CONTEXT *tcti_context, TSS2_SYS_CONTEXT *sys_context) {
    if (sys_context != NULL) {
        Tss2_Sys_Finalize(sys_context);
        free(sys_context);
    }
    if (tcti_context != NULL) {
        Tss2_TctiLdr_Finalize(&tcti_context);
    }
}
